# -*- coding: utf-8 -*-
"""
Author: chen xi (modified by Qwen)
Date: 2024/12/24 下午2:30
File: tcp_load_random.py
功能：高并发 TCP 客户端，发送随机字节心跳包，接收时不尝试 UTF-8 解码
"""

import asyncio
import random
import time

# -------------------------------
# 配置参数
# -------------------------------
HEARTBEAT_MIN_SIZE = 32    # 最小心跳包大小（字节）
HEARTBEAT_MAX_SIZE = 512   # 最大心跳包大小（字节）
# -------------------------------


async def connect_and_send_periodically(host, port, interval):
    try:
        reader, writer = await asyncio.open_connection(host, port)
        addr = writer.get_extra_info('peername')
        print(f"✅ 连接到 {addr}")

        async def send_data_task():
            while True:
                await asyncio.sleep(interval)

                # ✅ 生成随机长度的随机字节
                size = random.randint(HEARTBEAT_MIN_SIZE, HEARTBEAT_MAX_SIZE)
                random_bytes = bytes(random.getrandbits(8) for _ in range(size))

                try:
                    writer.write(random_bytes)
                    # 可选：打印发送日志（避免刷屏，可注释）
                    # print(f"📤 发送 {len(random_bytes)} 字节到 {addr}: {random_bytes.hex()[:32]}...")
                    await writer.drain()
                except Exception as e:
                    print(f"❌ 发送错误到 {addr}: {e}")
                    break

        # 启动发送任务
        asyncio.create_task(send_data_task())

        # 接收循环：只读原始 bytes，不 decode
        try:
            while True:
                data = await reader.read(1024)
                if not data:
                    break
                print(f"📥 从 {addr} 接收 {len(data)} 字节: {data.hex()}")

        except Exception as e:
            print(f"❌ 接收错误 from {addr}: {type(e).__name__}: {e}")
        finally:
            print(f"🔌 关闭连接 {addr}")
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()

    except Exception as e:
        print(f"❌ 连接失败 {host}:{port}: {e}")


async def create_connections(server_id, num_connections, host, port, interval):
    tasks = []

    for i in range(server_id * num_connections, (server_id + 1) * num_connections):
        # ✅ 每个连接独立任务，发送随机字节
        task = asyncio.create_task(
            connect_and_send_periodically(host, port, interval)
        )
        tasks.append(task)

        # 防止瞬时连接风暴
        if i % 1000 == 0 and i > 0:
            await asyncio.sleep(0.5)

    await asyncio.gather(*tasks)


if __name__ == '__main__':
    num_connections = 1000      # 连接数
    server_id = 0                # 服务器 ID（每台机器唯一）
    host = '120.26.136.5'       # 服务器地址
    port = 50000                 # 端口
    interval = 30               # 心跳间隔（秒）
    
    while True:
        # 运行客户端
        asyncio.run(create_connections(server_id, num_connections, host, port, interval))
        time.sleep(60)